{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ef12562-dcce-459d-907c-ed294ad11db0",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install modelscope ipywidgets tqdm accelerate diffusers transformers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20474daa",
   "metadata": {},
   "outputs": [],
   "source": [
    "from importlib.metadata import version\n",
    "print(\"diffusers version:\", version(\"diffusers\"))\n",
    "print(\"torch version:\", version(\"torch\"))\n",
    "print(\"transformers version:\", version(\"transformers\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16b0136b-ab36-41fd-9975-04a94dc30d4f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from model_download import download_model\n",
    "\n",
    "download_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "adfc6220",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "# Load the pipeline\n",
    "if torch.cuda.is_available():\n",
    "    torch_dtype = torch.bfloat16\n",
    "    device = \"cuda\"\n",
    "else:\n",
    "    torch_dtype = torch.float32\n",
    "    device = \"cpu\"\n",
    "print(f\"使用的是{device}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6dcdf673",
   "metadata": {},
   "outputs": [],
   "source": [
    "from diffusers import DiffusionPipeline\n",
    "\n",
    "model_name = \"./models/Qwen-Image\"\n",
    "\n",
    "pipe = DiffusionPipeline.from_pretrained(model_name, torch_dtype=torch_dtype,use_safetensors=True,device_map=\"balanced\")\n",
    "print(f\"管道已加载成功{pipe}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65566c8d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# pipe = pipe.to(device)\n",
    "\n",
    "positive_magic = {\n",
    "    \"en\": \", Ultra HD, 4K, cinematic composition.\", # for english prompt\n",
    "    \"zh\": \", 超清，4K，电影级构图.\" # for chinese prompt\n",
    "}\n",
    "\n",
    "# 定义5个测试提示词\n",
    "test_prompts = [\n",
    "    {\n",
    "        \"name\": \"cyberpunk_city\",\n",
    "        \"prompt\": '''A photo of a cyberpunk city at night, with neon lights reflecting on the rain-slicked streets. A ramen stall's sign clearly reads \"未来拉面\" (Future Ramen), with smaller text below saying \"AI-Powered Noodles.\" A nearby electronic billboard scrolls with binary code: \"01010100 01101111 01101110 01100111 01111001 01101001\".''',\n",
    "        \"description\": \"测试霓虹灯效果下的中英文文本渲染，以及复杂背景下数字/代码的清晰度\"\n",
    "    },\n",
    "    {\n",
    "        \"name\": \"ancient_book\",\n",
    "        \"prompt\": '''A close-up shot of a yellowed, ancient book open on a heavy wooden desk, with a quill pen beside it. On the page, beautiful calligraphy reads \"靜以修身，儉以養德\" (Cultivate oneself in tranquility, nurture virtue through thrift). At the bottom of the page, a faded English note says: \"The virtues are cultivated in tranquility\".''',\n",
    "        \"description\": \"测试不同中文书法风格和手写斜体英文的渲染，以及文本在羊皮纸等古老材质上的纹理\"\n",
    "    },\n",
    "    {\n",
    "        \"name\": \"product_packaging\",\n",
    "        \"prompt\": '''A high-end, close-up shot of a modern, minimalist skincare product box. The front of the box features the product name \"光子精华\" (Photon Essence) and the brand \"Quantum Glow.\" The ingredient list is clearly visible: \"Aqua, Glycerin, Peptide-5, Hyaluronic Acid (玻尿酸)\".''',\n",
    "        \"description\": \"测试现代产品设计上印刷文本的清晰度和纹理，评估模型处理中英混合排版、大小写敏感性和技术术语的能力\"\n",
    "    },\n",
    "    {\n",
    "        \"name\": \"street_graffiti\",\n",
    "        \"prompt\": '''A weathered brick wall in a city alley, covered in colorful graffiti art. At the center of the graffiti is a large Chinese character \"龙\" (Dragon), with \"DRAGON'S BREATH\" written next to it in a spray-paint style font. In the corner, a small signature reads: \"By Artist Qwen, 2024\".''',\n",
    "        \"description\": \"测试非标准艺术字体（如喷漆和手写）的渲染，以及文本与砖墙粗糙纹理的融合效果\"\n",
    "    },\n",
    "    {\n",
    "        \"name\": \"code_screen\",\n",
    "        \"prompt\": '''A programmer's desk with a computer monitor displaying code. A Python function is clearly visible on the screen: \"def calculate_fibonacci(n): # 计算斐波那契数列\\n    if n <= 1:\\n        return n\". A sticky note with the handwritten text \"Launch Deadline: 12/25\" is faintly visible in the screen's reflection.''',\n",
    "        \"description\": \"测试等宽字体（如代码字体）的准确渲染，包括缩进、语法高亮和双语注释，以及屏幕反射等复杂光学效果下的详细文本处理\"\n",
    "    }\n",
    "]\n",
    "\n",
    "negative_prompt = \" \" # Recommended if you don't use a negative prompt.\n",
    "\n",
    "# Generate with different aspect ratios\n",
    "aspect_ratios = {\n",
    "    \"1:1\": (1328, 1328),\n",
    "    \"16:9\": (1664, 928),\n",
    "    \"9:16\": (928, 1664),\n",
    "    \"4:3\": (1472, 1104),\n",
    "    \"3:4\": (1104, 1472),\n",
    "    \"3:2\": (1584, 1056),\n",
    "    \"2:3\": (1056, 1584),\n",
    "}\n",
    "\n",
    "width, height = aspect_ratios[\"16:9\"]\n",
    "\n",
    "# 依次生成5张图片\n",
    "for i, test_case in enumerate(test_prompts, 1):\n",
    "    print(f\"\\n正在生成第{i}张图片: {test_case['name']}\")\n",
    "    print(f\"测试目标: {test_case['description']}\")\n",
    "    print(f\"提示词: {test_case['prompt'][:100]}...\")\n",
    "    \n",
    "    image = pipe(\n",
    "        prompt=test_case['prompt'] + positive_magic[\"en\"],\n",
    "        negative_prompt=negative_prompt,\n",
    "        width=width,\n",
    "        height=height,\n",
    "        num_inference_steps=50,\n",
    "        true_cfg_scale=4.0,\n",
    "        generator=torch.Generator(device=\"cuda\").manual_seed(42 + i)  # 每张图片使用不同的种子\n",
    "    ).images[0]\n",
    "    \n",
    "    filename = f\"{test_case['name']}_test.png\"\n",
    "    image.save(filename)\n",
    "    print(f\"图片已保存为: {filename}\")\n",
    "    \n",
    "print(\"\\n所有5张测试图片生成完成！\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Qwen-Image (3.12.11)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
